import numpy

from model.wide_and_deep import WideAndDeep, WideAndDeep_plus
from model.loss import Loss
from torch_lib import traverse
from prework.metrics import accuracy, precision, recall, f1
from prework.preprocess import test_data
# from prework.preprocess_plus import test_data
import torch
from torch.utils.data import DataLoader, Dataset


class GetDataset(Dataset):

    def __init__(self):
        self.input = test_data()

    def __getitem__(self, index):
        x = self.input[index]
        for i in range(len(x)):
            x[i] = float(x[i])
        return torch.FloatTensor(x)

    def __len__(self):
        return len(self.input)


model = WideAndDeep()
# model = WideAndDeep_plus()
model.load_state_dict(torch.load('model/saved_model/model91.pt'))
dataset = GetDataset()
dataset = DataLoader(dataset, batch_size=1, shuffle=True)
loss = Loss()


def callback(data):
    y_pred = numpy.argmax(data['y_pred'])  # 模型预测
    print(y_pred)


traverse(model, dataset, callbacks=[callback], console_print=True)

